package encryption;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

import org.ietf.jgss.GSSContext;
import org.ietf.jgss.GSSException;
import org.ietf.jgss.MessageProp;

/**
 * This class, given a security context, handles data/messages encryption
 * @author pasquale
 * 
 */
public class EncryptionManager {

	private int BUFFER_SIZE = 1024;

	/**
	 * 
	 * @param context
	 * @param prop
	 * @param bufferSize size of receiving buffer, 0 for default (1024 bytes)
	 */
	public EncryptionManager(GSSContext context, MessageProp prop, int bufferSize){
		this.context = context;
		this.prop = prop;
		if(bufferSize != 0)
			BUFFER_SIZE = bufferSize;
	}

	public EncryptionManager(GSSContext context, MessageProp prop, int bufferSize,
			SecretKey key){
		this.encMode = true;
		this.context = context;
		this.prop = prop;
		try{
			byte[] iv = new byte[]{
					(byte)0x8E, 0x12, 0x39, (byte)0x9C,
					0x07, 0x72, 0x6F, 0x5A, 0x2C, 0x33,
					0x01, (byte)0x9A, 0x45, 0x12, (byte)0x89,
					0x04
			};
			AlgorithmParameterSpec paramSpec = new IvParameterSpec(iv);
			this.cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");	
			cipher.init(Cipher.ENCRYPT_MODE, key,paramSpec);
			this.decipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
			decipher.init(Cipher.DECRYPT_MODE, key,paramSpec);
		}catch(Exception ex){
			ex.printStackTrace();
			System.exit(-1);
		}


		if(bufferSize != 0)
			BUFFER_SIZE = bufferSize;
	}

	/**
	 * Listens on DataInputStream for encrypted messages
	 * @param context
	 * @param prop
	 * @return
	 * @throws IOException
	 * @throws GSSException
	 */
	public byte[] readEncMessage(DataInputStream inStream) throws IOException, GSSException{
		byte[] token = new byte[inStream.readInt()];
		inStream.readFully(token);
		byte[] bytes = context.unwrap(token, 0, token.length, prop);
		return bytes;
	}

	/**
	 * Writes an ecrypted message down to the Output Stream
	 * @param context
	 * @param prop
	 * @param message
	 * @throws GSSException
	 * @throws IOException
	 */
	public void writeEncMessage(byte[] message, DataOutputStream outStream) throws GSSException, IOException{
		byte[] token = context.wrap(message, 0, message.length, prop);
		outStream.writeInt(token.length);
		outStream.write(token);
		outStream.flush();
	}

	public void secureSend(InputStream in,OutputStream out) throws IOException, GSSException{
		if(encMode)
			in = new CipherInputStream(in,decipher);

		DataOutputStream dos = new DataOutputStream(out);
		byte[] buffer = new byte [BUFFER_SIZE];
		int nBytesRead = 0;
		while(nBytesRead >=0){
			nBytesRead = in.read(buffer, 0, buffer.length);
			if(nBytesRead > 0){
				byte token[] = context.wrap(buffer, 0, nBytesRead, prop);
				dos.writeInt(token.length);
				dos.write(token, 0, token.length);
				dos.flush();
			}
		}
		dos.writeInt(0);
	}


	/**
	 * 
	 * @param in encrypted stream
	 * @param out plain stream
	 * @throws IOException 
	 * @throws GSSException 
	 * 
	 */

	public void secureReceive(InputStream in,OutputStream out) throws IOException, GSSException{
		DataInputStream dis = new DataInputStream(in);
		if(encMode)
			out = new CipherOutputStream(out,cipher);

		byte[] buffer;
		int tokenL = 1;
		while(tokenL > 0){
			tokenL = dis.readInt();
			if(tokenL> 0){
				byte[] token=new byte[tokenL];
				dis.readFully(token);
				buffer = context.unwrap(token, 0, token.length, prop);
				out.write(buffer, 0, buffer.length);
			}
		}
		out.close();
	}

	public GSSContext getContext(){
		return context;
	}

	public MessageProp getMessageProps(){
		return prop;
	}


	private GSSContext context;
	private MessageProp prop;

	boolean encMode = false;
	private Cipher cipher;
	private Cipher decipher;

}
